#!/usr/bin/python
'''
Created on Feb 1, 2012

@author: xzhu
'''

import sys
from PROT import PROT

def load_key_res(ifile):
    """ load residue names from a file, such as "ASPA0085" """
    
    key_res = []
    lines = open(ifile).readlines()
    for line in lines:
        # "#" denotes a commented line
        if line.strip()[0] != '#':
            key_res.append(line.strip())

    return key_res


def extra_prot_info(prot):
    ''' 
    Display more information about the protein
    '''
    selfSum = 0.0
    pairSum = 0.0
    for conf in prot.conformers:
        if conf.confID[3:5] == 'DM': continue
        #if conf.focc < 0.01: continue
        selfSum += conf.selfEner * conf.focc
        #print "SELF:  confID: %s, focc: %.2f, selfEner: %.2f, selfEner * focc: %.2f, selfSum: %.2f" \
        #      % (conf.confID, conf.focc, conf.selfEner, conf.selfEner * conf.focc, selfSum)

        for oconf in prot.conformers:
            if oconf.confID[3:5] == 'DM': continue
            #if oconf.focc < 0.01: continue
            #if abs(conf.tPair[oconf.iConf]) < 0.1: continue
 
            pairSum += conf.tPair[oconf.iConf] * oconf.focc * conf.focc * 0.5
            #print "oconfID: %s, pairwise: %.2f, ofocc: %.2f, half pairwise * ofocc * focc: %.2f, pairSum: %.2f" \
            #      % (oconf.confID, conf.tPair[oconf.iConf], oconf.focc, conf.tPair[oconf.iConf] * oconf.focc * conf.focc, pairSum)

    print "\n\nAt last, selfSum: %.3f, pairSum: %.3f, total: %.3f" % (selfSum, pairSum, (selfSum + pairSum))
     
def main():
    
    # load all the info for the whole protein
    wholeProt = PROT()
    wholeProt.load_from_files()

    # get the fixed part and free part of the whole protein by identified key residues
    key_names = load_key_res(sys.argv[1])
    freeRes = []
    fixRes  = []
    
    for keyName in key_names:
        for res in wholeProt.residues:
            if res.resID == keyName:
                freeRes.append(res)
                break
            
    for res in wholeProt.residues:
        if res.resID not in key_names:
            fixRes.append(res)
            
    # initialize the free and fixed part from the residues already loaded
    freeProt = PROT()
    fixProt  = PROT()
    freeProt.init_from_res(freeRes)
    fixProt.init_from_res(fixRes)
    
    # calculate the self energy of the fixed part first, it doesn't change
    fixProt.get_self_ener()

    #print "For the fixed protein: "
    #extra_prot_info(fixProt)
    print "Total energy of fixed protein %.3f\n" % fixProt.total_ener
    
    # the header of the output
    #for res in freeRes:
    #    print '%-9s' % res.resID, 
    #print 'BoltzAve\tavg\tstd\tNms'
       
    freeProt.get_poss_iStates()
    freeProt.display_iStates(fixProt)
            
def helpMessage():
    print 'microEnergy.py f.txt'    

if __name__ == '__main__':
    '''
    calculate the free energy of different ionization states of a couple of residues
    from mfe
    '''
    if len(sys.argv) != 2:
        helpMessage()
    else:
        main()
